import cv2
import numpy as np
import matplotlib.pyplot as plt

# 1.使用同一相机，在两个不同的位姿分别拍摄同一个棋盘格，然后使用本质矩阵估计两次拍摄间的平移和旋转。
# 加载图片
img1 = cv2.imread('img_left.jpg', 0)
img2 = cv2.imread('img_right.jpg', 0)

# 初始化SIFT方法
sift = cv2.xfeatures2d_SIFT.create()

# 获取关键点和描述子
k1, d1 = sift.detectAndCompute(img1, None)
k2, d2 = sift.detectAndCompute(img2, None)

# 设置FLANN 超参数
FLANN_INDEX_KDTREE = 0
# K-D树索引超参数
index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
# 搜索超参数
search_params = dict(checks=50)

# 初始化FlannBasedMatcher匹配器
flann = cv2.FlannBasedMatcher(index_params, search_params)
# 通过KNN的方式匹配两张图的描述子
matches = flann.knnMatch(d1, d2, k=2)

good = []
pts1 = []
pts2 = []

# 筛选比较好的匹配点
for i, (m, n) in enumerate(matches):
    if m.distance < 0.8 * n.distance:
        good.append(m)
        pts2.append(k2[m.trainIdx].pt)
        pts1.append(k1[m.queryIdx].pt)

# 计算基础矩阵
pts1 = np.int32(pts1)
pts2 = np.int32(pts2)
# F为基本矩阵、mask是返回基本矩阵的值：没有找到矩阵，返回0，找到一个矩阵返回1，多个矩阵返回3
F, mask = cv2.findFundamentalMat(pts1, pts2, cv2.FM_LMEDS)
# 只选择有效数据
pts1 = pts1[mask.ravel() == 1]
pts2 = pts2[mask.ravel() == 1]


def drawlines(img1, img2, lines, pts1, pts2):
    """
    绘制图像极线
    :param img1:
    :param img2:
    :param lines:
    :param pts1:
    :param pts2:
    :return:
    """
    r, c = img1.shape
    img1 = cv2.cvtColor(img1, cv2.COLOR_GRAY2BGR)
    img2 = cv2.cvtColor(img2, cv2.COLOR_GRAY2BGR)

    for r, pt1, pt2 in zip(lines, pts1, pts2):
        color = tuple(np.random.randint(0, 255, 3).tolist())
        x0, y0 = map(int, [0, -r[2] / r[1]])
        x1, y1 = map(int, [c, -(r[2] + r[0] * c) / r[1]])

        img1 = cv2.line(img1, (x0, y0), (x1, y1), color, 1)
        img1 = cv2.circle(img1, tuple(pt1), 5, color, -1)
        img2 = cv2.circle(img2, tuple(pt2), 5, color, -1)
    return img1, img2


# 在右图（第二图）中找到与点相对应的极线，并在左图上画出它的线。
lines1 = cv2.computeCorrespondEpilines(pts2.reshape(-1, 1, 2), 2, F)
lines1 = lines1.reshape(-1, 3)
img5, img6 = drawlines(img1, img2, lines1, pts1, pts2)
# 找到与左图像（第一个图像）中的点对应的极线，以及在右图上画线
lines2 = cv2.computeCorrespondEpilines(pts1.reshape(-1, 1, 2), 1, F)
lines2 = lines2.reshape(-1, 3)
img3, img4 = drawlines(img2, img1, lines2, pts2, pts1)
plt.subplot(121), plt.imshow(img5)
plt.subplot(122), plt.imshow(img3)
plt.show()
##########################################
# 使用本质矩阵估计两次拍摄间的平移和旋转
E, masks = cv2.findEssentialMat(pts1, pts2, np.array([[4.70062852e+03, 0.00000000e+00, 1.38445962e+03],
 [0.00000000e+00, 4.72048901e+03, 2.28597865e+03],
 [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]], dtype=np.float32), cv2.FM_LMEDS)
print("本质矩阵", E)
# print("相对参数", masks)

print("..........................")
# SVD分解
print("开始矩阵分解")
# 创建矩阵A
A = np.array(E)

# 利用np.linalg.svd()函数直接进行奇异值分解
# 该函数有三个返回值：左奇异矩阵，所有奇异值，右奇异矩阵。
U, Sigma, VT = np.linalg.svd(A)

# 展示
print("旋转", U)
print("参数", Sigma)
print("平移", VT)

